package cn.demoncat.util.es.util;

import cn.demoncat.util.es.constant.EsConstant;
import cn.demoncat.util.lang.MapUtil;
import cn.demoncat.util.lang.StringUtil;
import cn.demoncat.util.lang.constant.StringConstant;
import cn.demoncat.util.lang.entity.Pager;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.domain.Sort.Direction;
import org.springframework.data.domain.Sort.NullHandling;
import org.springframework.data.domain.Sort.Order;
import org.springframework.data.elasticsearch.core.geo.GeoPoint;
import org.springframework.data.elasticsearch.core.query.*;

import java.util.*;
import java.util.Map.Entry;
import java.util.stream.Collectors;

/**
 * ElasticSearch Query工具
 * 
 * @see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html">官方文档</a>
 * 
 * @author 延晓磊
 *
 * @since 2020-9-2
 */
public class QueryUtil {

	/**
	 * 根据ES命令脚本查询
	 * 
	 * @param source	ES命令脚本
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static StringQuery build(String source) {
		return new StringQuery(source);
	}
	
	/**
	 * 根据ES原生API查询
	 * 
	 * @param builder	ES原生API
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 * 
	 * @see {@link QueryBuilders}
	 */
	public static NativeSearchQuery build(QueryBuilder builder) {
		return builder == null ? all() : new NativeSearchQuery(builder);
	}
	
	/**
	 * 根据ES原生API查询
	 * 
	 * @param builders	ES原生API：0个返回null，多个用bool.must组合
	 * @return 
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 * 
	 * @see {@link QueryBuilders}
	 */
	public static NativeSearchQuery build(List<QueryBuilder> builders) {
		if (CollectionUtils.isEmpty(builders)) {
			return all();
		}
		if (builders.size() == 1) {
			// 单个
			return new NativeSearchQuery(builders.get(0));
		}else {
			// 多个
			return new NativeSearchQuery(QueryBuilderUtil.and(builders));
		}
	}
	
	/**
	 * 根据ES原生API查询（复杂查询）
	 * 
	 * @param query		全文搜索，计分（如果非空，将只查询10000条，可以调用query.setTrackTotalHits(true)来取消优化）
	 * @param filter	精确过滤（缩小全文搜索的范围），不计分，不参与聚合运算
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 * 
	 * @see {@link QueryBuilders}
	 */
	public static NativeSearchQuery build(QueryBuilder query, QueryBuilder filter) {
		if (query == null && filter == null) {
			return all();
		}else if (query != null && filter == null) {
			return new NativeSearchQuery(query);
		}else if (query == null) {
			return new NativeSearchQuery(QueryBuilderUtil.all(), filter);
		}else{
			return new NativeSearchQuery(query, filter);
		}
	}
	
	/**
	 * 根据Criteria查询（简单查询）
	 * 
	 * @param criteria	查询条件
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static CriteriaQuery build(Criteria criteria) {
		return new CriteriaQuery(criteria);
	}
	
	/**
	 * 搜索全部命中
	 * 
	 * 1、为了加速搜索，ES采用了折衷方案（只搜索10000条），要搜索全部，需要在搜索时设置track_total_hits（调用本函数）
	 * 2、如果搜索10000+条时报错，需要设置索引的max_result_window（调用esAdminDao.setMaxResultWindow）
	 *
	 * @param query
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-11
	 * 
	 * @see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/7.9/search-your-data.html#track-total-hits">REST API</a>
	 */
	public static <T extends Query>T hitAll(T query) {
		query.setTrackTotalHits(true);
		return query;
	}
	
	/**
	 * 设置结果数
	 * 
	 * @param query		查询
	 * @param size		结果条数，默认返回10条
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static NativeSearchQuery addSize(NativeSearchQuery query, int size) {
		query.setMaxResults(size);
		return query;
	}
	
	/**
	 * 设置结果数
	 * 
	 * @param query		查询
	 * @param size		结果条数，默认返回10条
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static CriteriaQuery addSize(CriteriaQuery query, int size) {
		query.setMaxResults(size);
		return query;
	}
	
	/**
	 * 添加分页排序（搜索10000条）
	 * 
	 * @param query		查询
	 * @param page		分页排序
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addPage(T query, Pageable page) {
		return query.setPageable(page);
	}
	
	/**
	 * 添加分页排序（搜索全部命中）
	 * 
	 * @param query		查询
	 * @param page		分页排序
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addAllPage(T query, Pageable page) {
		return hitAll(query).setPageable(page);
	}
	
	/**
	 * 添加高亮
	 * 
	 * @param query		查询
	 * @param highlight	高亮条件
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addHigh(T query, HighlightQuery highlight) {
		if (highlight != null) {
			query.setHighlightQuery(highlight);
		}
		return query;
	}

	/**
	 * 添加高亮
	 *
	 * 注意：高亮字段，必须是分词+被搜索的字段
	 *
	 * @param query		查询
	 * @param fields	高亮字段
	 * @return
	 *
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addHigh(T query, String... fields) {
		return addHigh(query, getHigh(fields));
	}
	
	/**
	 * 添加聚合
	 * 
	 * 1个Query可以有n个Aggregation，1个Aggregation可以有n个Sub Aggregation
	 * 
	 * @param query
	 * @param aggs
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-4
	 */
	public static NativeSearchQuery addAgg(NativeSearchQuery query, AbstractAggregationBuilder<?>... aggs) {
		for (AbstractAggregationBuilder<?> agg : aggs) {
			query.addAggregation(agg);
		}
		return query;
	}

	/**
	 * 添加聚合
	 *
	 * 1个Query可以有n个Aggregation，1个Aggregation可以有n个Sub Aggregation
	 *
	 * @param query
	 * @param aggs
	 * @return
	 *
	 * @author 延晓磊
	 *
	 * @since 2020-9-4
	 */
	public static NativeSearchQuery addAgg(NativeSearchQuery query, Collection<AbstractAggregationBuilder<?>> aggs) {
		for (AbstractAggregationBuilder<?> agg : aggs) {
			query.addAggregation(agg);
		}
		return query;
	}
	
	/**
	 * 设置包含的字段
	 * 
	 * @param query	查询
	 * @param filed	结果中包含的字段（排除其它字段；只展示_id和field字段）
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addIncludes(T query, String... filed) {
		query.addSourceFilter(new FetchSourceFilter(filed, null));
		return query;
	}
	
	/**
	 * 设置排除的字段
	 * 
	 * @param query	查询
	 * @param filed	结果中排除的字段（展示其它字段；只排除field字段）
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static <T extends Query>T addExcludes(T query, String... filed) {
		query.addSourceFilter(new FetchSourceFilter(null, filed));
		return query;
	}
	
	/**
	 * 排除所有返回结果
	 * 
	 * 场景：只获取聚合结果
	 * 
	 * @param query
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-8
	 */
	public static NativeSearchQuery addExcludes(NativeSearchQuery query) {
		query.setMaxResults(0);
		query.addSourceFilter(new FetchSourceFilter(new String[]{EsConstant.ID_FIELD}, null));
		return query;
	}
	
	/**
	 * 排除所有返回结果
	 * 
	 * 场景：只获取聚合结果
	 * 
	 * @param query
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-8
	 */
	public static CriteriaQuery addExcludes(CriteriaQuery query) {
		query.setMaxResults(0);
		query.addSourceFilter(new FetchSourceFilter(new String[]{EsConstant.ID_FIELD}, null));
		return query;
	}
	
	/**
	 * 查询全部
	 * 
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-2
	 */
	public static NativeSearchQuery all() {
		return build(QueryBuilders.matchAllQuery());
	}
	
	/**
	 * 查询ID
	 * 
	 * @param id
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static CriteriaQuery id(Object id) {
		return build(CriteriaUtil.id(id));
	}
	
	/**
	 * 查询ID集
	 * 
	 * @param ids
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static CriteriaQuery ids(Collection<?> ids) {
		return build(CriteriaUtil.ids(ids));
	}
	
	/**
	 * 查询ID集
	 * 
	 * @param ids
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static NativeSearchQuery ids(String... ids) {
		return build(QueryBuilderUtil.ids(ids));
	}
	
	/**
	 * 域值查询（精确匹配）
	 * 
	 * @param field 域名
	 * @param value 域值
	 * 
	 * @return 
	 *
	 * @author 延晓磊
	 */
	public static CriteriaQuery kv(String field, Object value) {
		return build(CriteriaUtil.is(field, value));
	}

	/**
	 * 域值查询（分词搜索）
	 *
	 * @param field 域名
	 * @param value 域值
	 *
	 * @return
	 *
	 * @author 延晓磊
	 */
	public static NativeSearchQuery match(String field, String value) {
		return build(QueryBuilderUtil.match(field, value));
	}

	/**
	 * 域值对查询（精确匹配）
	 * 
	 * @param kvs 域值对
	 * 
	 * @return 
	 *
	 * @author 延晓磊
	 */
	public static CriteriaQuery kvs(Map<String, Object> kvs) {
		if (MapUtil.isEmpty(kvs)) {
			return build(CriteriaUtil.all());
		}
		return build(CriteriaUtil.is(kvs));
	}

	/**
	 * 域值对查询（分词搜索）
	 *
	 * @param kvs 域值对
	 *
	 * @return
	 *
	 * @author 延晓磊
	 */
	public static NativeSearchQuery matchs(Map<String, String> kvs) {
		if (MapUtil.isEmpty(kvs)) {
			return all();
		}
		return build(kvs.entrySet().stream().map(e -> QueryBuilderUtil.match(e.getKey(), e.getValue())).collect(Collectors.toList()));
	}
	
	/**
	 * 分页
	 * 
	 * @param pager
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static Pageable getPage(Pager<?> pager) {
		// 分页参数
		return PageRequest.of(pager.getPageNum() - 1, pager.getPageSize());
	}
	
	/**
	 * 分页排序
	 * 
	 * @param pager
	 * @param sort
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static Pageable getPage(Pager<?> pager, Sort sort) {
		if (sort == null) {
			return getPage(pager);
		}
		// 分页参数
		return PageRequest.of(pager.getPageNum() - 1, pager.getPageSize(), sort);
	}
	
	/**
	 * 排序方式
	 * 
	 * @param field 		排序的字段
	 * @param isAsc 		是否升序
	 * @param isNullLast 	是否Null排序在后
	 * 
	 * @return
	 */
	public static Order getOrder(String field, boolean isAsc, boolean isNullLast) {
		Direction direct = isAsc ? Direction.ASC : Direction.DESC;
		NullHandling nullHand = isNullLast ? NullHandling.NULLS_LAST : NullHandling.NULLS_FIRST;
		return new Order(direct, field, nullHand);
	}
	
	/**
	 * 排序参数
	 * 
	 * @param field
	 * @param isAsc
	 * @param isNullLast
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static Sort getSort(String field, boolean isAsc, boolean isNullLast) {
		return Sort.by(getOrder(field, isAsc, isNullLast));
	}
	
	/**
	 * 排序参数
	 * 
	 * @param field
	 * @param isAsc
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-3
	 */
	public static Sort getSort(String field, boolean isAsc) {
		return Sort.by(getOrder(field, isAsc, true));
	}
	
	/**
	 * 排序参数
	 * 
	 * @param sorts 排序方式 {字段名:0升序/1逆序}
	 *
	 * @return Sort.by(Sort.Order.desc("a"), Sort.Order.asc("b"))
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-1
	 */
	public static Sort getSort(LinkedHashMap<String, Integer> sorts) {
		if (MapUtil.isEmpty(sorts)) {
			return null;
		}
		// 转换
		List<Order> orders = new ArrayList<>();
		for (Entry<String, Integer> s : sorts.entrySet()) {
			orders.add(getOrder(s.getKey(), s.getValue() == 0, true));
		}
		// 转换为Sort
		return Sort.by(orders);
	}
	
	/**
	 * 
	 * 根据SQL创建排序参数
	 *
	 * @param sqlSort	SQL命令：xx ASC,yy DESC 或 xx,yy DESC格式
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2017年6月9日
	 */
	public static Sort getSort(String sqlSort) {
		if (StringUtils.isBlank(sqlSort)) {
			return null;
		}
		// 切分多个排序
		String[] sorts = StringUtil.splitStrict(sqlSort);
		String[] sort;
		List<Order> orders = new ArrayList<>();
		for (String s : sorts) {
			// 切分排序
			sort = s.trim().split(StringConstant.SPACING);
			orders.add(getOrder(sort[0], sort.length == 1 || sort[1].equals("ASC"), true));
		}
		return Sort.by(orders);
	}
	
	/**
	 * 地理空间排序
	 * 
	 * @param field
	 * @param point	基点
	 * 
	 * @return 最近优先
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-4
	 */
	public static Sort getSort(String field, GeoPoint point) {
		return Sort.by(new GeoDistanceOrder(field, point));
	}
	
	/**
	 * ID排序
	 * 
	 * @param isAsc
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-1
	 */
	public static Sort getSort(boolean isAsc) {
		Sort sort = Sort.by(EsConstant.ID_FIELD);
		return isAsc ? sort.ascending() : sort.descending();
	}
	
	/**
	 * 高亮
	 *
	 * 注意：高亮字段，必须是分词+被搜索的字段
	 * 
	 * @param fields
	 * @return
	 * 
	 * @author 延晓磊
	 *
	 * @since 2020-9-4
	 */
	public static HighlightQuery getHigh(String... fields) {
		HighlightBuilder builder = new HighlightBuilder();
		// 高亮字段
		for (String field : fields) {
			builder.field(field);
		}
		// 标签
		builder.preTags(EsConstant.HLT_PRE_TAG);
		builder.postTags(EsConstant.HLT_POST_TAG);
		return new HighlightQuery(builder);
	}
	
}
